import torch.onnx
from pytorch_model import CNNcar,Model, efficient_Regress

#Function to Convert to ONNX
def Convert_ONNX(model, input_size):

    # set the model to inference mode
    model.eval().cuda()
    device = torch.device('cuda')

    # Let's create a dummy input tensor
    dummy_input = torch.randn(1, *input_size, requires_grad=True, dtype=torch.float32).to(device)


    # Export the model
    torch.onnx.export(model,         # model being run
         dummy_input,       # model input (or a tuple for multiple inputs)
         "ImageRegress_10_31_resnet18.onnx",       # where to save the model
         export_params=True,  # store the trained parameter weights inside the model file
         opset_version=17,    # the ONNX version to export the model to
         do_constant_folding=True,  # whether to execute constant folding for optimization
         input_names = ['Image'],   # the model's input names
         output_names = ['X_Y_Z_vel'], # the model's output names
         dynamic_axes={'Image' : {0 : 'batch_size'},    # variable length axes
                                'X_Y_Z_vel' : {0 : 'batch_size'}})
    print(" ")
    print('Model has been converted to ONNX')


if __name__ == "__main__":
    # Let's build our model
    # train(5)
    # print('Finished Training')

    # Test which classes performed well
    # testAccuracy()

    # Let's load the model we just created and test the accuracy per label

    batch_size = 1  # 批处理大小
    input_shape = (1, 126, 224)  # 输入数据
    torch.set_default_tensor_type('torch.FloatTensor')
    # torch.set_default_tensor_type('torch.cuda.FloatTensor')
    model = Model(action_dim=3, max_action=1.0)
    # model = efficient_Regress(action_dim=3)
    path = "./trained_models/new-model-forward-201.pth"
    model.load_state_dict(torch.load(path))

    # Test with batch of images
    # testBatch()
    # Test how the classes performed
    # testClassess()

    # Conversion to ONNX
    Convert_ONNX(model, input_shape)
